import os
from pydantic import BaseModel
from pydantic_ai import Agent, Tool
from pydantic_ai.mcp import MCPServerSSE,MCPServerHTTP
from typing import Optional
from datetime import datetime, timedelta
from pydantic_ai.models.openai import OpenAIModel
from pydantic_ai.providers.openai import OpenAIProvider
import asyncio

# 创建MCP服务器连接 - 移除了不支持的retry参数
mcp_server = MCPServerSSE(url="http://127.0.0.1:8000/sse")

model = OpenAIModel(
    model_name="deepseek-chat",
    provider=OpenAIProvider(
        base_url='https://api.deepseek.com',
        api_key='sk-b7ad23e07ace4e96ab304ae01b45ae64'
        )
    )

async def main():
    # 创建Agent实例
    agent = Agent(model,mcp_servers=[mcp_server])

    # 启动 MCP 服务器连接
    async with agent.run_mcp_servers():
        print('MCP-客户端已启动')

        # 执行需要计算能力的查询
        result = await agent.run('明天天气怎样？你需要先使用工具获取明确的日期信息，然后再根据日期获取天气信息。')

    print(f"结果: {result.output}")

if __name__ == "__main__":
    asyncio.run(main())